package com.chenfan.common.extension.mybatisplus.enhanced;

import com.baomidou.mybatisplus.core.enums.SqlMethod;
import com.baomidou.mybatisplus.core.exceptions.MybatisPlusException;
import com.baomidou.mybatisplus.core.injector.AbstractMethod;
import com.baomidou.mybatisplus.core.injector.DefaultSqlInjector;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.extension.injector.methods.InsertBatchSomeColumn;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.keygen.NoKeyGenerator;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;

import java.util.List;

/**
 * 增强-默认SQL注入器
 *
 * @author wulg
 * @date 2021-09-13
 **/
public final class EnhancedSqlInjector extends DefaultSqlInjector {

    @Override
    public List<AbstractMethod> getMethodList (Class < ? > mapperClass){
        List<AbstractMethod> methodList = super.getMethodList(mapperClass);
        methodList.add(new InsertBatchSomeColumn() {
            @Override
            public String getMethod(SqlMethod sqlMethod) {
                // 自定义 mapper 方法名
                return "insertBatch";
            }
        });
        methodList.add(new InsertOrUpdateBath());
        methodList.add(new SelectOneForUpdate());
        methodList.add(new SelectListForUpdate());
        return methodList;
    }

    /**
     * @Description: 批量插入更新
     */
    public static class InsertOrUpdateBath extends AbstractMethod {

        //private static SqlMethod saveBatchSqlMethod = new SqlMethod("saveBatch", "批量保存数据", "<script>INSERT INTO %s %s VALUES %s ON DUPLICATE KEY UPDATE %s</script>");

        @Override
        public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
            if(!tableInfo.havePK()) {
                throw new MybatisPlusException("未标记主键, table:" + tableInfo.getTableName() + ", modelClass:" + modelClass.getName());
            }
            final SqlSource sqlSource = prepareSqlSource(tableInfo, modelClass);
            // 第三个参数必须和RootMapper的自定义方法名一致
            return this.addInsertMappedStatement(mapperClass, modelClass, "saveBatch", sqlSource, new NoKeyGenerator(), null, null);
        }

        /**
         * 拼接SQL语句
         * @param tableInfo
         * @param modelClass
         * @return
         */
        protected SqlSource prepareSqlSource(TableInfo tableInfo, Class<?> modelClass) {
            final String sql = "<script>INSERT INTO %s %s VALUES %s ON DUPLICATE KEY UPDATE %s</script>";
            final String tableName = tableInfo.getTableName();
            final String filedSql = prepareFieldSql(tableInfo);
            final String modelValuesSql = prepareModelValuesSql(tableInfo);
            final String duplicateKeySql = prepareDuplicateKeySql(tableInfo);
            final String sqlResult = String.format(sql, tableName, filedSql, modelValuesSql, duplicateKeySql);
            return languageDriver.createSqlSource(configuration, sqlResult, modelClass);
        }

        /**
         * 准备DuplicateKeySql语句
         * @param tableInfo
         * @return
         */
        protected String prepareDuplicateKeySql(TableInfo tableInfo) {
            StringBuilder duplicateKeySql = new StringBuilder();
            if (!StringUtils.isEmpty(tableInfo.getKeyColumn())) {
                duplicateKeySql.append(tableInfo.getKeyColumn()).append("=VALUES(").append(tableInfo.getKeyColumn()).append("),");
            }
            tableInfo.getFieldList().forEach(x -> {
                duplicateKeySql.append(x.getColumn()).append("=VALUES(").append(x.getColumn()).append("),");
            });
            duplicateKeySql.delete(duplicateKeySql.length() - 1, duplicateKeySql.length());
            return duplicateKeySql.toString();
        }

        /**
         * SQL语句拼接赋值
         * @param tableInfo
         * @return
         */
        protected String prepareModelValuesSql(TableInfo tableInfo) {
            StringBuilder valueSql = new StringBuilder();
            valueSql.append("<foreach collection=\"list\" item=\"item\" index=\"index\" open=\"(\" separator=\"),(\" close=\")\">");
            if (!StringUtils.isEmpty(tableInfo.getKeyProperty())) {
                valueSql.append("#{item.").append(tableInfo.getKeyProperty()).append("},");
            }
            tableInfo.getFieldList().forEach(x -> valueSql.append("#{item.").append(x.getProperty()).append("},"));
            valueSql.delete(valueSql.length() - 1, valueSql.length());
            valueSql.append("</foreach>");
            return valueSql.toString();
        }

        /**
         * 准备属性名
         **/
        protected String prepareFieldSql(TableInfo tableInfo) {
            StringBuilder fieldSql = new StringBuilder();
            fieldSql.append(tableInfo.getKeyColumn()).append(",");
            tableInfo.getFieldList().forEach(x -> {
                fieldSql.append(x.getColumn()).append(",");
            });
            fieldSql.delete(fieldSql.length() - 1, fieldSql.length());
            fieldSql.insert(0, "(");
            fieldSql.append(")");
            return fieldSql.toString();
        }
    }

    /**
     * 单个加悲观锁
     */
    public class SelectOneForUpdate extends AbstractMethod {

        private static final String MAPPER_METHOD = "selectOneForUpdate";
        private static final String SQL_TEMPLATE = "<script>%s SELECT %s FROM %s %s %s for update\n</script>";

        @Override
        public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
            SqlSource sqlSource = languageDriver.createSqlSource(configuration, String.format(SQL_TEMPLATE,
                    sqlFirst(), sqlSelectColumns(tableInfo, true), tableInfo.getTableName(),
                    sqlWhereEntityWrapper(true, tableInfo), sqlComment()), modelClass);
            return this.addSelectMappedStatementForTable(mapperClass, MAPPER_METHOD, sqlSource, tableInfo);
        }
    }

    /**
     * 批量加悲观锁
     */
    public class SelectListForUpdate extends AbstractMethod {

        private static final String MAPPER_METHOD = "selectListForUpdate";
        private static final String SQL_TEMPLATE = "<script>%s SELECT %s FROM %s %s %s for update\n</script>";

        @Override
        public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
            String sql = String.format(SQL_TEMPLATE, sqlFirst(), sqlSelectColumns(tableInfo, true), tableInfo.getTableName(),
                    sqlWhereEntityWrapper(true, tableInfo), sqlComment());
            SqlSource sqlSource = languageDriver.createSqlSource(configuration, sql, modelClass);
            return this.addSelectMappedStatementForTable(mapperClass, MAPPER_METHOD, sqlSource, tableInfo);
        }
    }


}
